import numpy as np
import numpy.random as npr
import matplotlib as mpl
import matplotlib.pyplot as plt
import copy, string, itertools, bisect, math, time


def build_allele_names(n):
    return ["a{}".format(i) for i in range(n)]
    
def build_phenotype_dict(alleles):
    return dict( zip(alleles,range(len(alleles))) )
    
# module wide variables, don't manually change here but use initialize function instead!
_FEMALE_FECUNDITY = 10
_NUMBER_OF_ALLELES = _FEMALE_FECUNDITY + 1
_ALLELES = build_allele_names(_NUMBER_OF_ALLELES)
_PHENOTYPE = build_phenotype_dict(_ALLELES)
_MK_LETHALITY = 1.
_MK_TRANSMISSION_RATE = 0.9
_REALLOCATION_EFFICIENCY = 0.2
_MUTATION_RATE = 1e-5
_POPULATION_CAPACITY = 100


def initialize(**pars):
    global _FEMALE_FECUNDITY, _NUMBER_OF_ALLELES, _ALLELES, _PHENOTYPE, _MK_LETHALITY, _MK_TRANSMISSION_RATE, _REALLOCATION_EFFICIENCY, _MUTATION_RATE, _POPULATION_CAPACITY
    globals().update(pars)
    _NUMBER_OF_ALLELES = _FEMALE_FECUNDITY + 1
    _ALLELES = build_allele_names(_NUMBER_OF_ALLELES)
    _PHENOTYPE = build_phenotype_dict(_ALLELES)
  
def randchoice(seq):
    return seq[npr.randint(len(seq))]
    
def random_genotype(alleles):
    return [randchoice(alleles), randchoice(alleles)]

def shift_mutate(a, alleles):
    i = alleles.index(a)
    if npr.random() < 0.5:                # shift left
        if not i == 0:
            i -= 1
    else:                               # shift right
        if not i == len(alleles)-1:
            i += 1
    return alleles[i]

def jump_mutate(a, alleles):
    mutant_alleles = alleles[:]
    mutant_alleles.remove(a)
    return randchoice( mutant_alleles )
        
def mutate(a, alleles, mode='shift'):
    if mode=='jump':
        return jump_mutate(a, alleles)
    else:
        return shift_mutate(a, alleles)

def fitness_compensation(x, F, b, rmin=None):
    """x    - number of killed sons
       F    - total clutch size
       rmin - resource threshold to viability
       b    - reallocation efficiency"""
    if not rmin:
        rmin = 1./(10*F)            # default value: optimal resource level is 10 times the threshold
    a = (1-F*rmin)*(F-x)
    numerator   = a + b*x
    denominator = a + b*x*F*rmin
    return numerator / denominator

def treeindent(s, pos):
    lines = s.splitlines()
    prefix = '    |'+' '*(pos-5)
    inline = lines[0]+'\n'
    for l in lines[1:]:
        inline += "{}{}\n".format(prefix,l)
    return inline.rstrip()

class WeightedRandomGenerator(object):
    def __init__(self, weights):
        self.totals = []
        running_total = 0

        for w in weights:
            running_total += w
            self.totals.append(running_total)

    def next(self):
        rnd = npr.random() * self.totals[-1]
        return bisect.bisect_right(self.totals, rnd)

    def __call__(self):
        return self.next()

class WeightedRandomGeneratorWithoutReplacement(object):
    def __init__(self, weights):
        self.totals = []
        self.selected = []
        running_total = 0

        for w in weights:
            running_total += w
            self.totals.append(running_total)

    def next(self):
        if len(self.totals) == len(self.selected):
            raise ValueError, "generator is empty"
        rnd = npr.random() * self.totals[-1]
        i = bisect.bisect_right(self.totals, rnd)
        if i not in self.selected:
            self.selected.append(i)
            return i
        else:
            return self.next()
        
    def __call__(self):
        return self.next()

class Malekiller(object):
    def __init__(self, lethality=_MK_LETHALITY, transmission_rate=_MK_TRANSMISSION_RATE, mk_id='A'):
        self.lethality = lethality
        self.transrate = transmission_rate
        self.id = mk_id
        
    def __str__(self):
        return """\
male killer ({self.id})
---------------
   |-- lethality: {self.lethality}
   |-- transmission rate: {self.transrate}""".format(self=self)
    
    def __repr__(self):
        return "Malekiller( lethality={self.lethality}, transmission_rate={self.transrate}, mk_id={self.id} )".format(self=self)

class Individual(object):
    def __init__(self, genotype, infection=None, sex='individual', fitness=1., _alleles=_ALLELES, _mutation_rate=_MUTATION_RATE):
        self.genotype = genotype
        self.infection = self.mk = infection
        self.infected = self.infection is not None             # bool value so you can test for infection with 'if self.infected:...'
        self.sex = sex
        self.fitness = fitness
        self._alleles = _alleles                                 # defaults to module variable
        self._phenodict = build_phenotype_dict(_alleles)
        self.phenotype = self.get_phenotypic_value()
        self._mutation_rate = _mutation_rate
        
    def __str__(self):
        return self.show()
        
    def __repr__(self):
        return "Individual( genotype={s.genotype}, infection={s.infection!r}, sex={s.sex}, fitness={s.fitness}, _alleles={s._alleles}, _mutation_rate={s._mutation_rate} )".format(s=self)
    
    def show(self):
        return """{self.sex}
{delimiter}
    |-- {infection}
    |-- phenotype: {self.phenotype}
    |-- fitness: {self.fitness}
    |-- genotype: {gtype}
    """.format( self=self, delimiter='-'*len(self.sex), infection=self.infection2string(), gtype=self.genotype2string() )
    
    def infection2string(self):
        if self.infected:
            return "infected with {}".format( treeindent(str(self.infection),22) )
        else:
            return "uninfected"
    
    def genotype2string(self):
        return '{g[0]} || {g[1]}'.format(g=self.genotype)
        
    def get_phenotypic_value(self):
        """Take the average."""
        a1,a2 = self.genotype
        return ( self._phenodict[a1] + self._phenodict[a2] ) / 2.
    
    def update_fitness(self, factor):
        self.fitness *= factor
    
    def gamete(self):
        mutation = 0
        g = randchoice( self.genotype )
        # mutation:
        if npr.random() < self._mutation_rate:
            g = mutate(g, self._alleles, mode='shift')
            mutation = 1
        return g, mutation

class Female(Individual):
    def __init__(self, genotype, infection=None, fitness=1., _alleles=_ALLELES, _fecundity=_FEMALE_FECUNDITY, _realloc_eff=_REALLOCATION_EFFICIENCY, _mutation_rate=_MUTATION_RATE):
        Individual.__init__(self, genotype, sex='female', infection=infection, fitness=fitness, _alleles=_alleles, _mutation_rate=_mutation_rate)
        self.sexratio = math.floor(self.phenotype)
        self.fecundity = _fecundity                       # this and the following default to module variable
        self._realloceff = _realloc_eff
    
    def __repr__(self):
        return "Female( genotype={s.genotype}, infection={s.infection!r}, fitness={s.fitness}, _alleles={s._alleles}, _fecundity={self._fecundity}, _realloc_eff=self._realloc_eff, _mutation_rate={s._mutation_rate} )".format(s=self)
    
    def show(self):
        return """{self.sex}
{delimiter}
    |-- {infection}
    |-- sex ratio: {self.sexratio}
    |-- fecundity: {self.fecundity}
    |-- fitness: {self.fitness}
    |-- genotype: {gtype}
    """.format( self=self, delimiter='-'*len(self.sex), infection=self.infection2string(), gtype=self.genotype2string() )
    
    def produce_brood_with(self, male):
        mut_count = 0
        daughters = []
        sons = []
        for i in range(1,self.fecundity+1):
            # offspring genotype:
            egg, m1   = self.gamete()
            sperm, m2 = male.gamete()
            mut_count += m1+m2
            # offspring infection:
            infection = None
            if self.infected:
                if npr.random() < self.mk.transrate:               # random transmission of infection
                    infection = self.mk
            # offspring sex:
            if i <= self.sexratio:
                sons.append( Male([egg,sperm], infection) )
            else:
                daughters.append( Female([egg,sperm], infection) )
        if self.infected:
            sons,daughters = self.male_killing(sons,daughters)
        male.number_of_matings += 1
        return sons, daughters, mut_count
    
    def male_killing(self, sons, daughters):
        n = len(sons)
        sons[:] = (s for s in sons if s.infected and not npr.random()<self.mk.lethality)
        body_count = n - len(sons)
        for offspring in sons+daughters:
            offspring.update_fitness( fitness_compensation(body_count, self.fecundity, self._realloceff, rmin=None) )
        return sons,daughters
    
class Male(Individual):
    def __init__(self, genotype, infection=None, fitness=1., _alleles=_ALLELES, _mutation_rate=_MUTATION_RATE):
        Individual.__init__(self, genotype, infection=infection, sex='male', fitness=fitness, _alleles=_alleles, _mutation_rate=_mutation_rate)
        self.number_of_matings = 0
        
    def show(self):
        return """{self.sex}
{delimiter}
    |-- {infection}
    |-- number of matings: {self.number_of_matings}
    |-- fitness: {self.fitness}
    |-- genotype: {gtype}
    """.format( self=self, delimiter='-'*len(self.sex), infection=self.infection2string(), gtype=self.genotype2string() )

    def __repr__(self):
        return "Male( genotype={s.genotype}, infection={s.infection!r}, fitness={s.fitness}, _alleles={s._alleles}, _mutation_rate={s._mutation_rate} )".format(s=self)
    

            
    
    
class Population(object):                             
    def __init__(self, females, males, capacity=_POPULATION_CAPACITY, generation=0, _alleles=_ALLELES, _female_fecundity=_FEMALE_FECUNDITY, _mk_lethality=_MK_LETHALITY, _mk_transmission_rate=_MK_TRANSMISSION_RATE, _realloc_eff=_REALLOCATION_EFFICIENCY, _mutation_rate=_MUTATION_RATE):
        self.females = females
        self.males = males
        self.capacity = capacity
        self.generation = generation
        self._alleles = _alleles
        self._female_fecundity = _female_fecundity
        self._reallocation_efficiency = _realloc_eff
        
        record_keys = ['generation','n_females','n_males','mk','mk_females','mk_males','n_mutations','average_sex_ratio','allele_count','genotype_count']
        self.record = {}
        for k in record_keys:
            self.record[k] = []
        self.update_record(mut_count=0)
    
    def allele_count(self, gcount=None):
        if gcount == None:
            gcount = self.genotype_count()
        return np.sum(gcount,0) + np.sum(gcount,1)
    
    def genotype_count(self):
        alleles = self._alleles
        nA = len(alleles)
        gcount = np.zeros((nA,nA), dtype=int)
        for i in self.females+self.males:
            a1,a2 = i.genotype
            gcount[alleles.index(a1),alleles.index(a2)] += 1
        return gcount
    
    def asr(self):
        """Calculate average sex ratio (asr) among females."""
        value = 0.
        for f in self.females:
            value += f.sexratio
        return value/len(self.females)
    
    def mk(self):
        nF = nM = 0
        for f in self.females:
            if f.infected:
                nF += 1
        for m in self.males:
            if m.infected:
                nM += 1
        return nM, nF
    
    def scan(self):
        """Basically, apply gcount, asr, and mk; but avoid iterating through population multiple times."""
        alleles = self._alleles
        nA = len(alleles)
        sr = 0.
        nF = nM = 0
        acount = [0 for a in alleles]           # allele count
        gcount = np.zeros((nA,nA), dtype=int)   # genotype count
        for f in self.females:
            sr += f.sexratio
            if f.infected:
                nF += 1
            a1,a2 = f.genotype
            gcount[ alleles.index(a1), alleles.index(a2) ] += 1
        sr /= len(self.females)
        for m in self.males:
            if m.infected:
                nM += 1
            a1,a2 = m.genotype
            gcount[ alleles.index(a1), alleles.index(a2) ] += 1
        acount = self.allele_count(gcount)
        return sr, nF, nM, acount, gcount            
    
    def update_record(self, mut_count):
        self.record['generation'].append( self.generation )
        self.record['n_males'].append( len(self.males) )
        self.record['n_females'].append( len(self.females) )
        self.record['n_mutations'].append( mut_count )
        sr, mkF, mkM, acount, gcount = self.scan()
        self.record['average_sex_ratio'].append( sr )
        self.record['mk_females'].append( mkF )
        self.record['mk_males'].append( mkM )
        self.record['mk'].append( mkF+mkM )
        self.record['allele_count'].append( acount )
        self.record['genotype_count'].append( gcount )
        
    def is_polymorphic(self):
        acount = self.record['allele_count'][-1]         # most current allele count
        maxcount = 2*self.capacity
        return maxcount not in acount
        
    def produce_offspring(self):
        mut_count = 0
        if self.males == []:
            raise ValueError, 'population extinct due to lack of males'
        males = self.males[:]
        npr.shuffle(males)
        females = self.females[:]
        npr.shuffle(females)
        temp_pool = []
        offspring = []

        for f in females:
            try:
                m = males.pop()
            except IndexError:
                males = temp_pool[:]
                npr.shuffle(males)
                temp_pool = []
                m = males.pop()
            sons,daughters,mut = f.produce_brood_with(m)
            mut_count += mut
            temp_pool.append(m)
            offspring.extend(sons+daughters)
        return offspring, mut_count
        
    def select_new_generation(self, offspring):
        females = []
        males = []
        weights = []
        for o in offspring:
            weights.append(o.fitness)
        wrg = WeightedRandomGeneratorWithoutReplacement(weights)
        for i in range(self.capacity):
            individual = offspring[ wrg() ]
            if individual.sex == 'female':
                females.append(individual)
            else:
                males.append(individual)
        return males, females
    
    def next(self):
        offspring, mut_count = self.produce_offspring()
        self.males, self.females = self.select_new_generation(offspring)
        self.generation += 1
        self.update_record(mut_count)
    
    def run(self, n, force=False, timing=False):
        """If force is False then n is maximal generation count but iteration stops if the population
        is no longer polymorphic."""
        if timing:
            g0 = self.generation
            t0 = time.time()
        if force:
            for i in range(n):
                self.next()
        else:
            g = 0.
            while self.is_polymorphic() and g<n:
                self.next()
                g += 1
        if timing:
            print "Iterated {} generations in {:.2f} seconds".format(self.generation-g0, time.time()-t0)
            
    def visualize(self, plots=['asr','nM','mk','mut','acount','inds']):
        F = self._female_fecundity
        labels = dict(asr=dict(rec='average_sex_ratio',xlabel='generation',ylabel='average sex ratio',style='k-',ylims=[0,F]), \
                      nM=dict(rec='n_males',xlabel='generation',ylabel='number of males',style='g-',ylims=[0,self.capacity]), \
                      mk=dict(rec='mk',xlabel='generation',ylabel='MK infected',style='r-',ylims=[0,self.capacity]), \
                      mut=dict(rec='n_mutations',xlabel='generation',ylabel='number of mutations',style='m-'), \
                      acount=dict(rec='allele_count',xlabel='generation',ylabel='alleles'), \
                      inds=dict(rec='individuals',xlabel='allele 1',ylabel='allele 2') )
        if not isinstance(plots, list):
            plots = list(plots)
        if 'inds' in plots:
            plots.remove('inds')
        figs = []
        g = np.array(self.record['generation'])
        for p in plots:
            f = plt.figure()
            ax = f.add_subplot(111)
            figs.append(f)
            rec = labels[p]['rec']
            xl = labels[p]['xlabel']
            yl = labels[p]['ylabel']
            
            a = np.array(self.record[rec])
            if p == 'mut':
                a = np.cumsum(a)
                labels[p]['ylims']= [0,np.max(a)]
            if p == 'acount':
                a = np.transpose( np.array(a,dtype=float) )
                a /= (2*self.capacity)                         # normalize (so that values range between 0 and 1)
                ma = np.ma.array(a,mask=(a==0))            # mask zero values
                cmap = mpl.cm.jet                          # use jet color scheme
                cmap.set_bad('w')                          # use white for zero values
                im = ax.imshow(ma,cmap=cmap,vmin=0.,vmax=1.,origin='lower',interpolation='nearest',aspect='auto')
                cb = plt.colorbar(im)
                cb.set_label('allele frequency')
            elif p == 'inds':
                pass
            else:
                ylims = labels[p]['ylims']
                s = labels[p]['style']
                ax.plot(g,a,s)
                ax.set_ylim(*ylims)
            ax.set_xlabel(xl)
            ax.set_ylabel(yl)
        return plots,figs
        
    def genotype_show(self, g=-1):
        f = plt.figure()
        ax = f.add_subplot(111)
        a = np.array( self.record['genotype_count'][g], dtype=float )
        a /= self.capacity                                      # normalize
        ma = np.ma.array(a, mask=(a==0))
        cmap = mpl.cm.jet
        cmap.set_bad('w')
        im = ax.imshow(ma,cmap=cmap,vmin=0.,vmax=1.,origin='lower',interpolation='nearest',aspect='auto')
        cb = plt.colorbar(im)
        cb.set_label('frequency')
        ax.set_xlabel('chromosome 1')
        ax.set_ylabel('chromosome 2')

def random_individual(mk=Malekiller()):
    alleles = _ALLELES
    if npr.random() < 0.5:
        infection = None
    else:
        infection = mk
    if npr.random() < 0.5:
        return Male(random_genotype(alleles), infection)
    else:
        return Female(random_genotype(alleles), infection)
        
def random_population(N, mk=Malekiller()):
    females = []
    males = []
    for i in range(N):
        individual = random_individual(mk)
        if individual.sex == 'male':
            males.append(individual)
        else:
            females.append(individual)
    return Population(females, males, capacity=N)

def uniform_population(gnum=1, mk=Malekiller()):
    """gnum - number of males/females per genotype"""
    alleles = _ALLELES
    nA = _NUMBER_OF_ALLELES
    females = []
    males = []
    for i in range(nA):
        for j in range(nA):
            for k in range(gnum):
                if npr.random() < 0.5:
                    infection = None
                else:
                    infection = mk
                male = Male([alleles[i],alleles[j]],infection)
                if npr.random() < 0.5:
                    infection = None
                else:
                    infection = mk
                female = Female([alleles[i],alleles[j]],infection)
                males.append(male)
                females.append(female)
    return Population(females,males,capacity=len(females+males))
        
    
if __name__ == '__main__':
#    npr.RandomState(21112011)
    npr.seed(21112011)
    
    parameters = dict(_FEMALE_FECUNDITY = 10, \
                      _MUTATION_RATE = 1e-5, \
                      _REALLOCATION_EFFICIENCY = 0.5, \
                      _MK_LETHALITY = 1., \
                      _MK_TRANSMISSION_RATE = 0.9)
    initialize(**parameters)
    
    
    pop = random_population(200)
    pop = uniform_population()
#    females = [Female(['a1','a2'],'uninfected') for i in range(30)]
#    males = [Male(['a1','a2'],'uninfected') for i in range(30)]
#    pop = Population(females,males)
    pop.run(10, force=True, timing=True)
    
    plots,figs = pop.visualize()    
    plt.show()
